import time

import cv2
import cvzone
from cvzone.SelfiSegmentationModule import SelfiSegmentation

segmentor = SelfiSegmentation()


# cap = cv2.VideoCapture(0)
# wi = 1280
# hi = 720
# cap.set(3, wi)
# cap.set(4, hi)
# cap.set(cv2.CAP_PROP_FPS, 60)
# segmentor = SelfiSegmentation()
# fpsReader = cvzone.FPS()
# imgBG = cv2.imread("background/background.jpg")
# print(imgBG.shape)
# imgBG = cv2.resize(imgBG, (wi, hi), cv2.IMREAD_UNCHANGED)
# print(imgBG.shape)
# while True:
#     success, img = cap.read()
#     imgOut = segmentor.removeBG(img, imgBG, threshold=0.1)
#     # debug
#     # imgStacked = cvzone.stackImages([img, imgOut], 2, 1)
#     # _, fps_iimgStackedmage = fpsReader.update(imgStacked, color=[0, 0, 255])
#
#     cv2.putText(imgOut, "1", (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
#                 2, (0, 0, 255), 2, cv2.LINE_AA)
#     # 照片 / 添加的文字 / 左上角坐标 / 字体 / 字体大小 / 颜色 / 字体粗细
#     cv2.imshow("image", imgOut)
#     cv2.imwrite("./results/" + "test.jpg", imgOut)
#     cv2.waitKey(1)
# start_time = time.time()
#         while True:
#             _, frame = cap.read()
#             frame = cv2.flip(frame, 1)
#             # cv2.putText(frame, className, (10, 50), cv2.FONT_HERSHEY_SIMPLEX,
#             #             1, (0, 0, 255), 2, cv2.LINE_AA)
#             frame = segmentor.removeBG(frame, imgBG, threshold=0.1)
#             end = time.time()-start_time
#             cv2.putText(frame, f"keep smile {int(end)}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX,
#                             1, (0, 0, 255), 2, cv2.LINE_AA)
#             cv2.imshow("Output", frame)
#             print(int(end))
#             if end >3:
#                 break
#     # Show the final output
#     cv2.imshow("Output", frame)


def getSmile(cap, background_pic, seconds, output_path, wi=1280, hi=720):
    cap.set(3, wi)
    cap.set(4, hi)
    imgBG = cv2.imread(background_pic)
    imgBG = cv2.resize(imgBG, (wi, hi), cv2.IMREAD_UNCHANGED)


    start_time = time.time()
    while True:
        _, frame = cap.read()
        frame = cv2.flip(frame, 1)
        frame = segmentor.removeBG(frame, imgBG, threshold=0.1)
        end = time.time() - start_time
        cv2.putText(frame, f"keep smile {int(end)}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX,
                                    1, (0, 0, 255), 2, cv2.LINE_AA)
        cv2.imshow("Output", frame)
        print(int(end))
        if end >3:
            break


if __name__ == '__main__':
    cap = cv2.VideoCapture(1)
    getSmile(cap, "background/background.jpg", 10, 'results/test.jpg')
